#include "OptimizationCostModel.h"
#include <iostream>
#include <algorithm>
#include <cmath>

// ==================== CostEvaluator 实现 ====================

CostEvaluator::CostEvaluator() {
    // 初始化优化权重
    optimizationWeights[OptimizationType::CONSTANT_FOLDING] = 1.0;
    optimizationWeights[OptimizationType::COMMON_SUBEXPRESSION_ELIMINATION] = 1.5;
    optimizationWeights[OptimizationType::DEAD_CODE_ELIMINATION] = 1.2;
    optimizationWeights[OptimizationType::ALGEBRAIC_OPTIMIZATION] = 1.3;
    optimizationWeights[OptimizationType::LOOP_OPTIMIZATION] = 2.0;
    optimizationWeights[OptimizationType::REGISTER_ALLOCATION] = 1.8;
    optimizationWeights[OptimizationType::INLINING] = 1.6;
    optimizationWeights[OptimizationType::VECTORIZATION] = 2.5;
}

OptimizationDecision CostEvaluator::evaluateOptimization(OptimizationType type, ASTNode* node, const OptimizationContext& context) {
    if (!node) {
        return OptimizationDecision(type, OptimizationBenefit::NONE, OptimizationCost::NONE, false, "Null node");
    }

    OptimizationBenefit benefit = OptimizationBenefit::NONE;
    OptimizationCost cost = OptimizationCost::NONE;
    bool shouldApply = false;
    std::string reason;

    switch (type) {
        case OptimizationType::CONSTANT_FOLDING:
            return evaluateConstantFolding(node, context);
        case OptimizationType::COMMON_SUBEXPRESSION_ELIMINATION:
            return evaluateCommonSubexpressionElimination(node, context);
        case OptimizationType::DEAD_CODE_ELIMINATION:
            return evaluateDeadCodeElimination(node, context);
        case OptimizationType::ALGEBRAIC_OPTIMIZATION:
            return evaluateAlgebraicOptimization(node, context);
        default:
            return OptimizationDecision(type, OptimizationBenefit::NONE, OptimizationCost::NONE, false, "Unsupported optimization");
    }
}

OptimizationDecision CostEvaluator::evaluateConstantFolding(ASTNode* node, const OptimizationContext& context) {
    if (!node) return OptimizationDecision(OptimizationType::CONSTANT_FOLDING, OptimizationBenefit::NONE, OptimizationCost::NONE, false, "Null node");

    // 检查是否可以进行常量折叠
    bool canFold = false;
    int complexity = calculateExpressionComplexity(node);

    if (node->type == "AddExpr" || node->type == "MulExpr" ||
        node->type == "RelExpr" || node->type == "LAndExpr" ||
        node->type == "LOrExpr") {

        // 检查子节点是否都是常量
        bool allConstants = true;
        for (const auto& child : node->children) {
            if (child->type != "NumLiteral") {
                allConstants = false;
                break;
            }
        }
        canFold = allConstants;
    }

    if (canFold) {
        OptimizationBenefit benefit = complexity > 5 ? OptimizationBenefit::HIGH :
                                     complexity > 2 ? OptimizationBenefit::MEDIUM : OptimizationBenefit::LOW;
        OptimizationCost cost = OptimizationCost::LOW;
        return OptimizationDecision(OptimizationType::CONSTANT_FOLDING, benefit, cost, true,
                                   "Constant expression can be folded, complexity: " + std::to_string(complexity));
    }

    return OptimizationDecision(OptimizationType::CONSTANT_FOLDING, OptimizationBenefit::NONE, OptimizationCost::NONE, false,
                               "Not a constant expression");
}

OptimizationDecision CostEvaluator::evaluateCommonSubexpressionElimination(ASTNode* node, const OptimizationContext& context) {
    if (!node) return OptimizationDecision(OptimizationType::COMMON_SUBEXPRESSION_ELIMINATION, OptimizationBenefit::NONE, OptimizationCost::NONE, false, "Null node");

    // 检查表达式在上下文中出现的频率
    std::string exprKey = generateExpressionKey(node);
    auto it = context.expressionFrequency.find(exprKey);

    if (it != context.expressionFrequency.end() && it->second > 1) {
        int frequency = it->second;
        int complexity = calculateExpressionComplexity(node);

        OptimizationBenefit benefit = frequency > 3 ? OptimizationBenefit::HIGH :
                                     frequency > 1 ? OptimizationBenefit::MEDIUM : OptimizationBenefit::LOW;
        OptimizationCost cost = complexity > 10 ? OptimizationCost::MEDIUM : OptimizationCost::LOW;

        return OptimizationDecision(OptimizationType::COMMON_SUBEXPRESSION_ELIMINATION, benefit, cost, true,
                                   "Expression appears " + std::to_string(frequency) + " times, complexity: " + std::to_string(complexity));
    }

    return OptimizationDecision(OptimizationType::COMMON_SUBEXPRESSION_ELIMINATION, OptimizationBenefit::NONE, OptimizationCost::NONE, false,
                               "No common subexpression found");
}

OptimizationDecision CostEvaluator::evaluateDeadCodeElimination(ASTNode* node, const OptimizationContext& context) {
    if (!node) return OptimizationDecision(OptimizationType::DEAD_CODE_ELIMINATION, OptimizationBenefit::NONE, OptimizationCost::NONE, false, "Null node");

    bool isDeadCode = false;
    std::string reason;

    // 检查是否是未使用的变量声明
    if (node->type == "DeclAssignStmt") {
        std::string varName = extractVariableName(node);
        auto it = context.variableUsage.find(varName);
        if (it == context.variableUsage.end() || it->second == 0) {
            isDeadCode = true;
            reason = "Unused variable: " + varName;
        }
    }

    // 检查是否是无副作用的表达式语句
    if (node->type == "Expr" && !hasSideEffects(node)) {
        isDeadCode = true;
        reason = "Expression statement without side effects";
    }

    if (isDeadCode) {
        return OptimizationDecision(OptimizationType::DEAD_CODE_ELIMINATION, OptimizationBenefit::MEDIUM, OptimizationCost::LOW, true, reason);
    }

    return OptimizationDecision(OptimizationType::DEAD_CODE_ELIMINATION, OptimizationBenefit::NONE, OptimizationCost::NONE, false, "Not dead code");
}

OptimizationDecision CostEvaluator::evaluateAlgebraicOptimization(ASTNode* node, const OptimizationContext& context) {
    if (!node) return OptimizationDecision(OptimizationType::ALGEBRAIC_OPTIMIZATION, OptimizationBenefit::NONE, OptimizationCost::NONE, false, "Null node");

    bool canOptimize = false;
    std::string reason;

    // 检查代数恒等式
    if (node->type == "AddExpr" || node->type == "MulExpr" ||
        node->type == "LAndExpr" || node->type == "LOrExpr") {

        if (node->children.size() == 2) {
            ASTNode* left = node->children[0].get();
            ASTNode* right = node->children[1].get();

            // 检查 x + 0, x * 1, x && true 等恒等式
            if ((node->type == "AddExpr" && isZeroConstant(right)) ||
                (node->type == "MulExpr" && isOneConstant(right)) ||
                (node->type == "LAndExpr" && isTrueConstant(right))) {
                canOptimize = true;
                reason = "Algebraic identity optimization";
            }
        }
    }

    if (canOptimize) {
        return OptimizationDecision(OptimizationType::ALGEBRAIC_OPTIMIZATION, OptimizationBenefit::LOW, OptimizationCost::NONE, true, reason);
    }

    return OptimizationDecision(OptimizationType::ALGEBRAIC_OPTIMIZATION, OptimizationBenefit::NONE, OptimizationCost::NONE, false, "No algebraic optimization possible");
}

CodeQualityMetrics CostEvaluator::evaluateCodeQuality(ASTNode* node) {
    CodeQualityMetrics metrics;

    if (!node) return metrics;

    // 递归计算代码质量指标
    calculateMetricsRecursive(node, metrics);

    return metrics;
}

void CostEvaluator::calculateMetricsRecursive(ASTNode* node, CodeQualityMetrics& metrics) {
    if (!node) return;

    // 计算指令数
    if (isInstructionNode(node)) {
        metrics.instructionCount++;
    }

    // 计算内存访问
    if (node->type == "VarRef" || node->type == "AssignStmt") {
        metrics.memoryAccessCount++;
    }

    // 计算控制流复杂度
    if (node->type == "IfStmt" || node->type == "WhileStmt" || node->type == "ForStmt") {
        metrics.controlFlowComplexity++;
    }

    // 递归处理子节点
    for (const auto& child : node->children) {
        calculateMetricsRecursive(child.get(), metrics);
    }
}

int CostEvaluator::calculateExpressionComplexity(ASTNode* node) {
    if (!node) return 0;

    int complexity = 1; // 基础复杂度

    // 根据节点类型增加复杂度
    if (node->type == "AddExpr" || node->type == "MulExpr") {
        complexity += 2;
    } else if (node->type == "RelExpr" || node->type == "LAndExpr" || node->type == "LOrExpr") {
        complexity += 3;
    } else if (node->type == "FuncCall") {
        complexity += 5;
    }

    // 递归计算子节点复杂度
    for (const auto& child : node->children) {
        complexity += calculateExpressionComplexity(child.get());
    }

    return complexity;
}

// ==================== OptimizationCoordinator 实现 ====================

OptimizationCoordinator::OptimizationCoordinator() : maxIterations(10), convergenceThreshold(0.01) {
    // 设置默认优化顺序
    optimizationOrder = {
        OptimizationType::CONSTANT_FOLDING,
        OptimizationType::ALGEBRAIC_OPTIMIZATION,
        OptimizationType::COMMON_SUBEXPRESSION_ELIMINATION,
        OptimizationType::DEAD_CODE_ELIMINATION,
        OptimizationType::LOOP_OPTIMIZATION,
        OptimizationType::REGISTER_ALLOCATION,
        OptimizationType::INLINING,
        OptimizationType::VECTORIZATION
    };

    // 启用所有优化
    for (auto type : optimizationOrder) {
        optimizationEnabled[type] = true;
    }
}

void OptimizationCoordinator::coordinateOptimizations(ASTNode* root) {
    if (!root) return;

    std::cerr << "Starting coordinated optimization..." << std::endl;

    // 迭代优化直到收敛
    bool converged = optimizeUntilConvergence(root);

    if (converged) {
        std::cerr << "Optimization converged successfully." << std::endl;
    } else {
        std::cerr << "Optimization reached maximum iterations." << std::endl;
    }
}

bool OptimizationCoordinator::optimizeUntilConvergence(ASTNode* root) {
    CodeQualityMetrics previousMetrics = costEvaluator.evaluateCodeQuality(root);

    for (int iteration = 0; iteration < maxIterations; iteration++) {
        std::cerr << "Optimization iteration " << (iteration + 1) << "/" << maxIterations << std::endl;

        // 为每个节点选择最优优化
        optimizeNodeRecursively(root);

        // 评估当前代码质量
        CodeQualityMetrics currentMetrics = costEvaluator.evaluateCodeQuality(root);

        // 检查是否收敛
        if (hasConverged(previousMetrics, currentMetrics)) {
            std::cerr << "Optimization converged at iteration " << (iteration + 1) << std::endl;
            return true;
        }

        previousMetrics = currentMetrics;
    }

    return false;
}

void OptimizationCoordinator::optimizeNodeRecursively(ASTNode* node) {
    if (!node) return;

    // 创建优化上下文
    OptimizationContext context(node);
    context.currentMetrics = costEvaluator.evaluateCodeQuality(node);

    // 选择最优优化
    auto decisions = selectOptimalOptimizations(node, context);

    // 应用优化决策
    applyOptimizationDecisions(node, decisions);

    // 递归处理子节点
    for (auto& child : node->children) {
        optimizeNodeRecursively(child.get());
    }
}

std::vector<OptimizationDecision> OptimizationCoordinator::selectOptimalOptimizations(ASTNode* node, const OptimizationContext& context) {
    std::vector<OptimizationDecision> decisions;

    // 评估所有可用的优化
    for (auto type : optimizationOrder) {
        if (optimizationEnabled[type]) {
            auto decision = costEvaluator.evaluateOptimization(type, node, context);
            if (decision.shouldApply) {
                decisions.push_back(decision);
            }
        }
    }

    // 按净收益排序
    std::sort(decisions.begin(), decisions.end(),
              [](const OptimizationDecision& a, const OptimizationDecision& b) {
                  return a.netBenefit > b.netBenefit;
              });

    return decisions;
}

void OptimizationCoordinator::applyOptimizationDecisions(ASTNode* root, const std::vector<OptimizationDecision>& decisions) {
    for (const auto& decision : decisions) {
        std::cerr << "Applying optimization: " << static_cast<int>(decision.type)
                  << " (benefit: " << static_cast<int>(decision.benefit)
                  << ", cost: " << static_cast<int>(decision.cost)
                  << ", net: " << decision.netBenefit << ")" << std::endl;

        // 这里应该调用相应的优化器
        // 实际实现需要与现有的优化器集成
    }
}

bool OptimizationCoordinator::hasConverged(const CodeQualityMetrics& before, const CodeQualityMetrics& after) {
    double improvement = std::abs(after.estimatedExecutionTime - before.estimatedExecutionTime) /
                        (before.estimatedExecutionTime + 1e-6);

    return improvement < convergenceThreshold;
}

// 辅助函数实现
std::string CostEvaluator::generateExpressionKey(ASTNode* node) {
    if (!node) return "";

    std::string key = node->type;
    if (!node->value.empty()) {
        key += ":" + node->value;
    }

    for (const auto& child : node->children) {
        key += "(" + generateExpressionKey(child.get()) + ")";
    }

    return key;
}

std::string CostEvaluator::extractVariableName(ASTNode* node) {
    if (!node || node->type != "DeclAssignStmt") return "";

    for (const auto& child : node->children) {
        if (child->type == "ID") {
            return child->value;
        }
    }

    return "";
}

bool CostEvaluator::hasSideEffects(ASTNode* node) {
    if (!node) return false;

    if (node->type == "FuncCall" || node->type == "AssignStmt") {
        return true;
    }

    for (const auto& child : node->children) {
        if (hasSideEffects(child.get())) {
            return true;
        }
    }

    return false;
}

bool CostEvaluator::isZeroConstant(ASTNode* node) {
    return node && node->type == "NumLiteral" && node->value == "0";
}

bool CostEvaluator::isOneConstant(ASTNode* node) {
    return node && node->type == "NumLiteral" && node->value == "1";
}

bool CostEvaluator::isTrueConstant(ASTNode* node) {
    return node && node->type == "NumLiteral" && node->value == "1";
}

bool CostEvaluator::isInstructionNode(ASTNode* node) {
    return node && (node->type == "AddExpr" || node->type == "MulExpr" ||
                   node->type == "AssignStmt" || node->type == "FuncCall");
}

// 添加缺失的 OptimizationCoordinator 方法实现
void OptimizationCoordinator::setOptimizationOrder(const std::vector<OptimizationType>& order) {
    optimizationOrder = order;
}

void OptimizationCoordinator::enableOptimization(OptimizationType type, bool enabled) {
    optimizationEnabled[type] = enabled;
}

void OptimizationCoordinator::setMaxIterations(int iterations) {
    maxIterations = iterations;
}

void OptimizationCoordinator::setConvergenceThreshold(double threshold) {
    convergenceThreshold = threshold;
}

void OptimizationCoordinator::optimizeFunction(ASTNode* funcNode) {
    if (!funcNode || funcNode->type != "FuncDef") return;

    // 查找函数体
    for (auto& child : funcNode->children) {
        if (child->type == "Block") {
            optimizeBlock(child.get());
            break;
        }
    }
}

void OptimizationCoordinator::optimizeBlock(ASTNode* blockNode) {
    if (!blockNode || blockNode->type != "Block") return;

    std::vector<std::unique_ptr<ASTNode>> optimizedStatements;

    for (auto& child : blockNode->children) {
        if (child->type == "ExprStmt") {
            optimizeExpressionStatement(child.get(), optimizedStatements);
        } else {
            optimizedStatements.push_back(std::move(child));
        }
    }

    // 替换原始语句
    blockNode->children.clear();
    for (auto& stmt : optimizedStatements) {
        blockNode->children.push_back(std::move(stmt));
    }
}

void OptimizationCoordinator::optimizeExpressionStatement(ASTNode* exprStmt, std::vector<std::unique_ptr<ASTNode>>& optimizedStatements) {
    if (!exprStmt || exprStmt->type != "ExprStmt") return;

    // 这里可以添加表达式语句的优化逻辑
    // 目前只是简单地保留原始语句
    optimizedStatements.push_back(std::unique_ptr<ASTNode>(new ASTNode(*exprStmt)));
}

void OptimizationCoordinator::generateOptimizationReport(const std::vector<OptimizationDecision>& decisions) {
    std::cerr << "\n=== Optimization Report ===" << std::endl;
    std::cerr << "Total optimizations applied: " << decisions.size() << std::endl;

    for (const auto& decision : decisions) {
        std::cerr << "Optimization: " << static_cast<int>(decision.type)
                  << " - Benefit: " << static_cast<int>(decision.benefit)
                  << ", Cost: " << static_cast<int>(decision.cost)
                  << ", Net: " << decision.netBenefit
                  << " - " << decision.reason << std::endl;
    }
}
